#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
绘制修正偏差后的控制数据图表
处理*_data_bais.csv文件，生成与control_yu_improved.pdf格式一致的图表

使用方法:
python plot_bias_corrected_data.py --data_dir _results/threetanks/4/control_fig
"""

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import argparse
import os
from pathlib import Path
import glob

def load_bias_corrected_data(data_dir):
    """加载修正偏差后的CSV数据文件"""
    data = {}
    
    # 查找所有*_data_bais.csv文件
    bias_files = glob.glob(os.path.join(data_dir, '*_data_bais_LD.csv'))
    
    if not bias_files:
        print(f"错误: 在目录 {data_dir} 中未找到*_data_bais_LD.csv文件")
        return None
    
    print(f"找到 {len(bias_files)} 个偏差修正数据文件:")
    
    for file_path in bias_files:
        file_name = os.path.basename(file_path)
        print(f"  - {file_name}")
        
        try:
            df = pd.read_csv(file_path)
            data[file_name] = df
        except Exception as e:
            print(f"加载 {file_name} 失败: {e}")
    
    return data

def create_bias_corrected_plot(data, output_dir):
    """创建修正偏差后的控制图表"""
    
    # 分离y和u数据
    y_data = {}
    u_data = {}
    
    for file_name, df in data.items():
        if file_name.startswith('y') and 'bais' in file_name:
            # 提取输出编号 (y1, y2, y3)
            output_num = file_name.split('_')[0]
            y_data[output_num] = df
        elif file_name.startswith('u') and 'bais' in file_name:
            # 提取输入编号 (u1, u2, u3)
            input_num = file_name.split('_')[0]
            u_data[input_num] = df
    
    n_outputs = len(y_data)
    n_inputs = len(u_data)
    
    if n_outputs == 0 and n_inputs == 0:
        print("错误: 未找到有效的y或u数据")
        return False
    
    print(f"输出数量: {n_outputs}, 输入数量: {n_inputs}")
    
    # 创建图表布局
    n_cols = max(n_outputs, n_inputs)
    fig, ax = plt.subplots(nrows=2, ncols=n_cols, sharex=True, figsize=(16, 8))
    
    # 确保ax是2D数组
    if n_cols == 1:
        ax = ax.reshape(2, 1)
    
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 绘制输出数据 (第一行)
    for i in range(n_cols):
        if i < n_outputs:
            output_key = f'y{i+1}'
            if output_key in y_data:
                df = y_data[output_key]
                time = df['Time']
                
                # 绘制DDeePC数据 (使用红色实线，与原图一致)
                if 'LDeePC' in df.columns:
                    ax[0, i].plot(time, df['LDeePC'], '-', color='red', linewidth=1, label='LDeePC mean')
                if 'DDeePC' in df.columns:
                    ax[0, i].plot(time, df['DDeePC'], '-', color='purple', linewidth=1, label='DDeePC mean')
                # 绘制DeePC数据 (使用蓝色点划线，与原图一致)
                if 'DeePC' in df.columns:
                    deepc_data = df['DeePC'].dropna()
                    if not deepc_data.empty:
                        deepc_time = time[df['DeePC'].notna()]
                        ax[0, i].plot(deepc_time, deepc_data, '-.', color='dodgerblue', linewidth=1, label='DeePC mean')
                
                # 绘制开环数据 (使用橙色点划线，与原图一致)
                if 'Open_loop' in df.columns:
                    ax[0, i].plot(time, df['Open_loop'], '-.', color='darkorange', linewidth=1, label='Open-loop mean')
                
                # 绘制设定点 (使用绿色虚线，与原图一致)
                if 'Setpoint' in df.columns:
                    ax[0, i].plot(time, df['Setpoint'], '--', color='g', linewidth=1, label='Set-point')
                
                ax[0, i].set_ylabel(f'y{i+1}', fontsize=12)
                ax[0, i].grid(True, alpha=0.3)
                
                # 自适应y轴范围
                y_min = min(df[col].min() for col in ['DDeePC', 'Open_loop', 'Setpoint'] if col in df.columns and df[col].notna().any())
                y_max = max(df[col].max() for col in ['DDeePC', 'Open_loop', 'Setpoint'] if col in df.columns and df[col].notna().any())
                if 'DeePC' in df.columns and df['DeePC'].notna().any():
                    y_min = min(y_min, df['DeePC'].min())
                    y_max = max(y_max, df['DeePC'].max())
                
                y_range = y_max - y_min
                ax[0, i].set_ylim(y_min - 0.1 * y_range, y_max + 0.1 * y_range)
        else:
            ax[0, i].set_visible(False)
    
    # 绘制输入数据 (第二行)
    for i in range(n_cols):
        if i < n_inputs:
            input_key = f'u{i+1}'
            if input_key in u_data:
                df = u_data[input_key]
                time = df['Time']
                
                # 绘制DDeePC数据 (使用红色实线，与原图一致)
                if 'LDeePC' in df.columns:
                    ax[1, i].plot(time, df['LDeePC'], '-', color='red', linewidth=1, label='LDeePC mean')
                if 'DDeePC' in df.columns:
                    ax[1, i].plot(time, df['DDeePC'], '-', color='purple', linewidth=1, label='DDeePC mean')
                # 绘制设定点 (使用绿色虚线，与原图一致)
                if 'Setpoint' in df.columns:
                    ax[1, i].plot(time, df['Setpoint'], '--', color='g', linewidth=1, label='Set-point')
                
                ax[1, i].set_ylabel(f'u{i+1}', fontsize=12)
                ax[1, i].set_xlabel('Time (day)', fontsize=12)
                ax[1, i].grid(True, alpha=0.3)
                
                # 自适应y轴范围
                y_min = df['DDeePC'].min()
                y_max = df['DDeePC'].max()
                if 'Setpoint' in df.columns:
                    y_min = min(y_min, df['Setpoint'].min())
                    y_max = max(y_max, df['Setpoint'].max())
                
                y_range = y_max - y_min
                ax[1, i].set_ylim(y_min - 0.1 * y_range, y_max + 0.1 * y_range)
        else:
            ax[1, i].set_visible(False)
    
    # 添加图例（只在第一个子图添加，与原图一致）
    if n_outputs > 0:
        ax[0, 0].legend(loc='best', fontsize=10)
    
    # 设置整体标题（简化标题，与原图风格一致）
    fig.suptitle('修正偏差后的控制性能对比', fontsize=14)
    
    # 调整布局
    plt.tight_layout()
    plt.subplots_adjust(top=0.93)
    
    # 保存图表
    output_path = os.path.join(output_dir, 'control_bias_corrected.pdf')
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"图表已保存至: {output_path}")
    
    # 也保存PNG格式
    png_path = os.path.join(output_dir, 'control_bias_corrected.png')
    plt.savefig(png_path, dpi=300, bbox_inches='tight')
    print(f"图表已保存至: {png_path}")
    
    plt.show()
    
    return True

def batch_process_directories(base_dir):
    """批量处理多个实验目录"""
    processed_count = 0
    
    # 查找所有包含control_fig目录的实验目录
    for root, dirs, files in os.walk(base_dir):
        if 'control_fig' in dirs:
            control_fig_path = os.path.join(root, 'control_fig')
            
            # 检查是否有偏差修正数据文件
            bias_files = glob.glob(os.path.join(control_fig_path, '*_data_bais.csv'))
            if bias_files:
                print(f"\n处理目录: {control_fig_path}")
                
                # 加载并绘制数据
                data = load_bias_corrected_data(control_fig_path)
                if data:
                    success = create_bias_corrected_plot(data, control_fig_path)
                    if success:
                        processed_count += 1
                        print(f"✅ 成功处理: {control_fig_path}")
                    else:
                        print(f"❌ 处理失败: {control_fig_path}")
    
    return processed_count

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='绘制修正偏差后的控制数据图表')
    parser.add_argument('--data_dir', type=str, required=True,
                        help='包含*_data_bais.csv文件的目录路径')
    parser.add_argument('--output_dir', type=str, default=None,
                        help='输出目录（默认与数据目录相同）')
    parser.add_argument('--batch', action='store_true',
                        help='批量处理模式，递归处理所有子目录中的control_fig文件夹')
    
    args = parser.parse_args()
    
    # 检查数据目录是否存在
    if not os.path.exists(args.data_dir):
        print(f"错误: 数据目录不存在: {args.data_dir}")
        return
    
    if args.batch:
        # 批量处理模式
        print(f"批量处理模式，搜索目录: {args.data_dir}")
        processed_count = batch_process_directories(args.data_dir)
        print(f"\n批量处理完成！共处理了 {processed_count} 个目录")
    else:
        # 单个目录处理模式
        # 设置输出目录
        output_dir = args.output_dir if args.output_dir else args.data_dir
        os.makedirs(output_dir, exist_ok=True)
        
        print(f"数据目录: {args.data_dir}")
        print(f"输出目录: {output_dir}")
        
        # 加载数据
        data = load_bias_corrected_data(args.data_dir)
        if data is None:
            return
        
        # 创建图表
        success = create_bias_corrected_plot(data, output_dir)
        
        if success:
            print("\n✅ 修正偏差后的控制图表生成完成！")
        else:
            print("\n❌ 图表生成失败")

if __name__ == '__main__':
    main()